# memento/models.py
from django.db import models
from django.contrib.auth.models import User
from django.db.models import JSONField

from algorithm.models import UserTagMastery
from algorithm.models.user_problem_recommendation import UserProblemRecommendation, UserComprehensiveLevel
from recommendation.models import SolvedProblem


class UserStateMemento(models.Model):
    user = models.OneToOneField(User, on_delete=models.CASCADE, primary_key=True)
    comprehensive_level_data = models.JSONField(default=dict)
    problems_recommendation_data = models.JSONField(default=list)
    tag_mastery_data = models.JSONField(default=list)
    solved_problems_data = models.JSONField(default=list)

    @staticmethod
    def create_memento(user_id):
        user = User.objects.get(pk=user_id)  # 主键，这里pk可换为id
        memento, created = UserStateMemento.objects.update_or_create(
            user=user,
            defaults={
                # 默认的反索引是user.类名_set
                # 也可以自定义related_name
                'comprehensive_level_data': user.usercomprehensivelevel.save_state(),
                'problems_recommendation_data': [rec.save_state() for rec in user.userproblemrecommendation_set.all()],
                'tag_mastery_data': [tag.save_state() for tag in user.usertagmastery_set.all()],
                'solved_problems_data': [sp.save_state() for sp in user.solved_problems.all()],
            }
        )
        return memento

    @staticmethod
    def restore_memento(user_id):
        user = User.objects.get(pk=user_id)
        memento = UserStateMemento.objects.get(user=user)
        UserComprehensiveLevel.restore_state(user, memento.comprehensive_level_data)
        # OneToOne的特殊用法
        UserProblemRecommendation.restore_state(user, memento.problems_recommendation_data)
        UserTagMastery.restore_state(user, memento.tag_mastery_data)
        SolvedProblem.restore_state(user, memento.solved_problems_data)
